﻿//-----------------------------------------------------------------------------------------------------------
// WBFSSync Project by Omega Frost 
// http://wbfssync.codeplex.com/
//
// WBFSSync is Licensed under the terms of the 
// Microsoft Reciprocal License (Ms-RL)
//-----------------------------------------------------------------------------------------------------------
using System;
using System.IO;
using System.Runtime.InteropServices;
using System.Text;
using Microsoft.Win32.SafeHandles;
using WBFSe3.IO;

namespace WBFSe3.Native
{
    public static class WbfsNative
    {
        //------------------------------------------------------------------------------------------
        // Constants
        //------------------------------------------------------------------------------------------
        
        public static readonly IntPtr InvalidHandle = new IntPtr(-1);


        //------------------------------------------------------------------------------------------
        // Creates a direct-opening string
        //------------------------------------------------------------------------------------------
        internal static string CreateDirectOpenString(Char drive)
        {
            return @"\\.\" + drive + ':';
        }


        //------------------------------------------------------------------------------------------
        // Opens a drive for raw read/write
        //------------------------------------------------------------------------------------------
        internal static IntPtr OpenDrive(String drive)
        {
            return OpenDrive(drive[0], FileAccess.ReadWrite);
        }


        //------------------------------------------------------------------------------------------
        // Opens a drive for raw read/write
        //------------------------------------------------------------------------------------------
        internal static IntPtr OpenDrive(Char drive)
        {
            return OpenDrive(drive, FileAccess.ReadWrite);
        }


        //------------------------------------------------------------------------------------------
        // Opens a drive with the desired access
        //------------------------------------------------------------------------------------------
        internal static IntPtr OpenDrive(Char drive, FileAccess access)
        {
            String sDrive = CreateDirectOpenString(drive);
            IntPtr pDrive = CreateFile(sDrive, access,
                FileShare.ReadWrite, IntPtr.Zero, FileMode.Open,
                EFileAttributes.NoBuffering | EFileAttributes.Write_Through, IntPtr.Zero);

            if (pDrive == InvalidHandle)
            {
                int error = Marshal.GetLastWin32Error();
                switch (error)
                {
                    default: throw new WbfsIOException(error);
                    case 2: throw new FileNotFoundException();
                    case 3: throw new DirectoryNotFoundException();
                    case 5: throw new UnauthorizedAccessException();
                    case 32: throw new SharingViolationException();
                }
            }

            return pDrive;
        }


        //------------------------------------------------------------------------------------------
        // Locks a volume
        //------------------------------------------------------------------------------------------
        internal static void LockVolume(IntPtr handle)
        {
            uint dummy = 0;
            if (DeviceIoControl(handle, EIOControlCode.FsctlLockVolume, IntPtr.Zero, 0,
                IntPtr.Zero, 0, ref dummy, IntPtr.Zero) == 0)
                throw new WbfsIOException();
        }


        //------------------------------------------------------------------------------------------
        // Unlocks a volume
        //------------------------------------------------------------------------------------------
        internal static void UnlockVolume(IntPtr handle)
        {
            uint dummy = 0;
            if (DeviceIoControl(handle, EIOControlCode.FsctlUnlockVolume, IntPtr.Zero, 0,
                IntPtr.Zero, 0, ref dummy, IntPtr.Zero) == 0)
                throw new WbfsIOException();
        }


        //------------------------------------------------------------------------------------------
        // Dismounts a volume
        //------------------------------------------------------------------------------------------
        internal static void DismountVolume(IntPtr handle)
        {
            uint dummy = 0;
            if (DeviceIoControl(handle, EIOControlCode.FsctlDismountVolume, IntPtr.Zero, 0,
                IntPtr.Zero, 0, ref dummy, IntPtr.Zero) == 0)
                throw new WbfsIOException();
        }


        //------------------------------------------------------------------------------------------
        // Gets the physical disk's geometry
        //------------------------------------------------------------------------------------------
        internal static SDISKGEOMETRY GetDiskGeometry(IntPtr handle)
        {
            uint bytes = 0;
            SDISKGEOMETRY dg = new SDISKGEOMETRY();
            IntPtr pdg = Marshal.AllocHGlobal(Marshal.SizeOf(dg));

            int r = DeviceIoControl(handle, EIOControlCode.DiskGetDriveGeometry, 
                IntPtr.Zero, 0, pdg, (uint)Marshal.SizeOf(dg), ref bytes, IntPtr.Zero);

            dg = (SDISKGEOMETRY)Marshal.PtrToStructure(pdg, typeof(SDISKGEOMETRY));
            Marshal.FreeHGlobal(pdg);

            if (r == 0) throw new WbfsIOException();
            return dg;
        }


        //------------------------------------------------------------------------------------------
        // Gets the disk's partition information
        //------------------------------------------------------------------------------------------
        internal static SPARTITIONINFO GetPartitionInfo(IntPtr handle)
        {
            uint bytes = 0;
            SPARTITIONINFO pi = new SPARTITIONINFO();
            IntPtr ppi = Marshal.AllocHGlobal(Marshal.SizeOf(pi));

            int r = DeviceIoControl(handle, EIOControlCode.DiskGetPartitionInfo,
                IntPtr.Zero, 0, ppi, (uint)Marshal.SizeOf(pi), ref bytes, IntPtr.Zero);

            pi = (SPARTITIONINFO)Marshal.PtrToStructure(ppi, typeof(SPARTITIONINFO));
            Marshal.FreeHGlobal(ppi);

            if (r == 0) throw new WbfsIOException();
            return pi;
        }


        //------------------------------------------------------------------------------------------
        // Creates a file handle
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        internal static extern IntPtr CreateFile(String lpFileName, 
            [MarshalAs(UnmanagedType.U4)]FileAccess dwDesiredAccess,
            [MarshalAs(UnmanagedType.U4)]FileShare dwShareMode, 
            IntPtr lpSecurityAttributes,
            [MarshalAs(UnmanagedType.U4)]FileMode dwCreationDisposition,
            [MarshalAs(UnmanagedType.I4)]EFileAttributes dwFlagsAndAttributes,
            IntPtr hTemplateFile);


        //------------------------------------------------------------------------------------------
        // Reads data from a file
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool ReadFile(IntPtr hFile, IntPtr lpBuffer, uint nNumberOfBytesToRead,
            out uint lpNumberOfBytesRead, IntPtr lpOverlapped);


        //------------------------------------------------------------------------------------------
        // Reads data from a file
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool ReadFile(IntPtr hFile, [Out] byte[] lpBuffer, uint nNumberOfBytesToRead,
            out uint lpNumberOfBytesRead, IntPtr lpOverlapped);


        //------------------------------------------------------------------------------------------
        // Writes data to a file
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        internal static extern bool WriteFile(IntPtr hFile, IntPtr lpBuffer, uint nNumberOfBytesToWrite,
            out uint lpNumberOfBytesWritten, IntPtr lpOverlapped);


        //------------------------------------------------------------------------------------------
        // Writes data to a file
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        internal static extern bool WriteFile(IntPtr hFile, [In] byte[] lpBuffer, uint nNumberOfBytesToWrite,
            out uint lpNumberOfBytesWritten, IntPtr lpOverlapped);


        //------------------------------------------------------------------------------------------
        // Sets the position in a file
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        internal static extern bool SetFilePointerEx(IntPtr hFile, long liDistanceToMove,
            out long lpNewFilePointer, SeekOrigin dwMoveMethod);


        //------------------------------------------------------------------------------------------
        // Sets the end of a file
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern bool SetEndOfFile(IntPtr hFile);


        //------------------------------------------------------------------------------------------
        // Flush the file buffers
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        internal static extern bool FlushFileBuffers(IntPtr hFile);


        //------------------------------------------------------------------------------------------
        // Closes a handle
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        internal static extern bool CloseHandle(IntPtr hObject);


        //------------------------------------------------------------------------------------------
        // Retrieves information about a volume
        //------------------------------------------------------------------------------------------
        [DllImport("Kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
        internal extern static bool GetVolumeInformation(string RootPathName, 
            StringBuilder VolumeNameBuffer, int VolumeNameSize, 
            out uint VolumeSerialNumber, out uint MaximumComponentLength,
            out EFileSystemFeature FileSystemFlags, StringBuilder FileSystemNameBuffer,
            int nFileSystemNameSize);


        //------------------------------------------------------------------------------------------
        // Retrieves the ammount of free space of a disk
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        [return: MarshalAs(UnmanagedType.Bool)]
        internal static extern bool GetDiskFreeSpaceEx(string lpDirectoryName,
            out ulong lpFreeBytesAvailable, out ulong lpTotalNumberOfBytes,
            out ulong lpTotalNumberOfFreeBytes);


        //------------------------------------------------------------------------------------------
        // Get the drive type
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true)]
        internal static extern uint GetDriveType([In, MarshalAs(UnmanagedType.LPStr)] 
            string lpRootPathName);


        //------------------------------------------------------------------------------------------
        // Set the drive label
        //------------------------------------------------------------------------------------------
        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        [return: MarshalAs(UnmanagedType.Bool)]
        internal static extern bool SetVolumeLabel(
            [In, MarshalAs(UnmanagedType.LPStr)] string lpDirectoryName,
            [In, MarshalAs(UnmanagedType.LPStr)] string lpVolumeName);


        //------------------------------------------------------------------------------------------
        // Sends a control command to the device
        //------------------------------------------------------------------------------------------
        [DllImport("Kernel32.dll", SetLastError = false, CharSet = CharSet.Auto)]
        internal static extern int DeviceIoControl(IntPtr hDevice,
            EIOControlCode IoControlCode, [In] IntPtr InBuffer, uint nInBufferSize, 
            [Out] IntPtr OutBuffer, uint nOutBufferSize, ref uint pBytesReturned, 
            [In] IntPtr Overlapped);
    }
}
